

package mat

import (
	"ry/kernel"
)

type matrix struct {
	iOffset int
	jOffset int
	rows    int
	cols    int
	stride  int
	backed  Matrix
	data    []float64
}

// only methods which need access to the internal data structure
// should be in the interface; all others should be operators/functions
type Matrix interface {
	Grid
	At (i,j int) float64
	Set (i,j int, v float64)
}

func Zeros (rows,cols int) Matrix {

	m := new (matrix)

	m.rows = rows
	m.cols = cols
	m.stride = cols
	m.backed = m
	m.data = make ([]float64, rows*cols)

	return m
}

func (m *matrix) Index (i,j int) int {

	rows, cols := m.backed.Dims ()

	return (m.iOffset + i) % rows * m.stride + (m.jOffset + j) % cols
}

func (m *matrix) At (i,j int) float64 {
	return m.data [m.Index(i,j)]
}

func (m *matrix) Set(i,j int, v float64) {

	m.data [m.Index(i,j)] = v
}

func (m *matrix) NumElems () int {

	return m.rows * m.cols
}

func (m *matrix) Dims () (int, int) {

	return m.rows, m.cols
}

func (m *matrix) WrapIndex (i,j int) (int,int) {

	for i < 0 { i += m.rows }
	for i >= m.rows { i -= m.rows }
	for j < 0 { j += m.cols }
	for j >= m.cols { j -= m.cols } 

	return i,j
}

func (m *matrix) Iter () <-chan float64 {

	rows, cols := m.Dims ()

	next := make (chan float64)

	go func () {
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				next <- m.At (i,j)
			}
		}
		close (next)
	} ()
	
	return next
}

func (m *matrix) Select (i1,j1, i2,j2 int) Grid {

	i1,j1 = m.WrapIndex (i1,j1)
	i2,j2 = m.WrapIndex (i2,j2)

	r := *m

	if i2 < i1 {	// vertical wrap
		r.rows = m.rows - i1 + i2 + 1
	} else {
		r.rows = i2 - i1 + 1
	}

	if j2 < j1 {	// horizontal wrap
		r.cols = m.cols - j1 + j2 + 1
	} else {
		r.cols = j2 - j1 + 1
	}

	r.iOffset = m.iOffset + i1
	r.jOffset = m.jOffset + j1

	return &r
}

func (in *matrix) Convolve (out Matrix, r int, k func (window Matrix, center *float64)) {

	rows, cols := in.Dims ()

	kernel.Launch (rows, cols, 
		func (i,j int) bool {
			temp := out.At (i,j)
			k (in.Select (i-r,j-r, i+r,j+r).(Matrix), &temp)
			out.Set (i,j, temp)
			return true
		})
}


